# Copyright 2023 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

"""Train Cell."""
import math
import mindspore as ms
import mindspore.ops as ops
import mindspore.nn as nn
from mindspore import context
from mindspore.ops import operations as P
from mindspore.ops import functional as F
from mindspore.ops import composite as C
from mindspore.common.tensor import Tensor
from mindspore.common import dtype as mstype
from mindspore.common.parameter import Parameter
from mindspore.communication.management import get_group_size
from mindspore.nn.wrap.grad_reducer import DistributedGradReducer
from mindspore.context import ParallelMode
from mindspore.train.serialization import load_checkpoint, load_param_into_net
from .tinybert_model import BertModelCLS


GRADIENT_CLIP_TYPE = 1
GRADIENT_CLIP_VALUE = 1.0

_grad_overflow = C.MultitypeFuncGraph("_grad_overflow")
grad_overflow = P.FloatStatus()
clip_grad = C.MultitypeFuncGraph("clip_grad")
@clip_grad.register("Number", "Number", "Tensor")
def _clip_grad(clip_type, clip_value, grad):
    """
    Clip gradients.

    Inputs:
        clip_type (int): The way to clip, 0 for 'value', 1 for 'norm'.
        clip_value (float): Specifies how much to clip.
        grad (tuple[Tensor]): Gradients.

    Outputs:
        tuple[Tensor], clipped gradients.
    """
    if clip_type not in (0, 1):
        return grad
    dt = F.dtype(grad)
    if clip_type == 0:
        new_grad = C.clip_by_value(grad, F.cast(F.tuple_to_array((-clip_value,)), dt),
                                   F.cast(F.tuple_to_array((clip_value,)), dt))
    else:
        new_grad = nn.ClipByNorm()(grad, F.cast(F.tuple_to_array((clip_value,)), dt))
    return new_grad

grad_scale = C.MultitypeFuncGraph("grad_scale")
reciprocal = P.Reciprocal()

@grad_scale.register("Tensor", "Tensor")
def tensor_grad_scale(scale, grad):
    return grad * reciprocal(scale)

@_grad_overflow.register("Tensor")
def _tensor_grad_overflow(grad):
    return grad_overflow(grad)


def compute_neuron_head_importance(base_model, base_config, new_param_dict, eval_dataset):
    # load model
    base_model.set_train(False)
    # prepare things for heads
    n_layers, n_heads = base_config.num_hidden_layers, base_config.num_attention_heads
    head_importance = ops.zeros((n_layers, n_heads), ms.float32)
    head_mask = ops.ones((n_layers, n_heads), ms.float32)
    # collect weights
    intermediate_weight = []
    intermediate_bias = []
    output_weight = []

    for key, value in new_param_dict.items():
        if 'intermediate' in key:
            if value.dim() > 1:
                intermediate_weight.append(value)
            else:
                intermediate_bias.append(value)

        if 'output' in key and 'attention' not in key:
            if value.dim() > 1:
                output_weight.append(value)

    neuron_importance = []
    for w in intermediate_weight:
        neuron_importance.append(ops.Ones()((w.shape[0]), ms.float32))

    columns_list = ["input_ids", "input_mask", "segment_ids", "label_ids"]
    for step, data in enumerate(eval_dataset.create_dict_iterator()):
        input_data = []
        for i in columns_list:
            input_data.append(data[i])
        input_ids, input_mask, token_type_id, label_ids = input_data
        grad_fn = ms.grad(base_model, grad_position=3)
        grad_head_mask = grad_fn(input_ids, token_type_id, input_mask, head_mask)

        head_importance += ops.abs(grad_head_mask)

        # calculate  neuron importance
        len_w = len(intermediate_weight)
        selected_w1 = ()
        selected_b1 = ()
        selected_w2 = ()
        current_importance = ()
        for i in range(len_w):
            selected_w1 += (intermediate_weight[i],)
            selected_b1 += (intermediate_bias[i],)
            selected_w2 += (output_weight[i],)
            current_importance += (neuron_importance[i],)
        for i in range(len_w):
            grad_wb = ms.grad(base_model, grad_position=None, weights=(selected_w1[i], selected_b1[i], selected_w2[i]))
            wb_grad = grad_wb(input_ids, token_type_id, input_mask, head_mask)
            w1_grad = wb_grad[0]
            b1_grad = wb_grad[1]
            w2_grad = wb_grad[2]
            current_importance += ops.abs((selected_w1[i] * w1_grad).sum(axis=1) + (selected_b1[i] * b1_grad))
            current_importance += ops.abs((selected_w2[i] * w2_grad).sum(axis=0))

    return base_model, head_importance, current_importance


def reorder_neuron_head(base_model, head_importance, neuron_importance):
    """ reorder neurons based on their importance.
        Arguments:
            model: bert model
            head_importance: 12*12 matrix for head importance in 12 layers
            neuron_importance: list for neuron importance in 12 layers.
    """
    # reorder heads and ffn neurons
    for layer, current_importance in enumerate(neuron_importance):
        # reorder heads
        idx = ops.Sort(descending=True)(head_importance[layer])[-1]
        base_model.bert.bert_encoder.layers[layer].attention.reorder_heads(idx)
        # reorder neurons
        idx = ops.Sort(descending=True)(current_importance)[-1]
        base_model.bert.bert_encoder.layers[layer].intermediate.reorder_neurons(idx)
        base_model.bert.bert_encoder.layers[layer].output.reorder_neurons(idx)


class ClipGradients(nn.Cell):
    """
    Clip gradients.

    Args:
        grads (list): List of gradient tuples.
        clip_type (Tensor): The way to clip, 'value' or 'norm'.
        clip_value (Tensor): Specifies how much to clip.

    Returns:
        List, a list of clipped_grad tuples.
    """
    def __init__(self):
        super(ClipGradients, self).__init__()
        self.clip_by_norm = nn.ClipByNorm()
        self.cast = P.Cast()
        self.dtype = P.DType()

    def construct(self,
                  grads,
                  clip_type,
                  clip_value):
        """clip gradients"""
        if clip_type not in (0, 1):
            return grads
        new_grads = ()
        for grad in grads:
            dt = self.dtype(grad)
            if clip_type == 0:
                t = C.clip_by_value(grad, self.cast(F.tuple_to_array((-clip_value,)), dt),
                                    self.cast(F.tuple_to_array((clip_value,)), dt))
            else:
                t = self.clip_by_norm(grad, self.cast(F.tuple_to_array((clip_value,)), dt))
            new_grads = new_grads + (t,)
        return new_grads


class SoftmaxCrossEntropy(nn.Cell):
    """SoftmaxCrossEntropy loss"""
    def __init__(self):
        super(SoftmaxCrossEntropy, self).__init__()
        self.log_softmax = P.LogSoftmax(axis=-1)
        self.softmax = P.Softmax(axis=-1)
        self.reduce_mean = P.ReduceMean()
        self.cast = P.Cast()

    def construct(self, predicts, targets):
        likelihood = self.log_softmax(predicts)
        target_prob = self.softmax(targets)
        loss = self.reduce_mean(-target_prob * likelihood)

        return self.cast(loss, mstype.float32)


class BertNetworkWithLoss(nn.Cell):
    """
    Provide bert pre-training loss through network.
    Args:
        teacher_config (BertConfig): The config of BertModel.
        is_training (bool): Specifies whether to use the training mode.
        use_one_hot_embeddings (bool): Specifies whether to use one-hot for embeddings. Default: False.
    Returns:
        Tensor, the loss of the network.
    """
    def __init__(self, eval_dataset, teacher_config, teacher_ckpt, student_config, student_ckpt,
                 is_training, task_type, num_labels, use_one_hot_embeddings=False,
                 temperature=1.0, dropout_prob=0.1, is_rep_fit=True,
                 training_phase='dynabertw'):
        super(BertNetworkWithLoss, self).__init__()
        # load teacher model
        self.teacher = BertModelCLS(teacher_config, False, num_labels, dropout_prob,
                                    use_one_hot_embeddings)
        param_dict = load_checkpoint(teacher_ckpt)
        new_param_dict = {}
        for key, value in param_dict.items():
            new_key = 'teacher.' + key
            new_param_dict[new_key] = value
        load_param_into_net(self.teacher, new_param_dict)

        # no_grad
        self.teacher.set_train(False)
        params = self.teacher.trainable_params()
        for param in params:
            param.requires_grad = False

        # load student model
        self.bert = BertModelCLS(student_config, is_training, num_labels, dropout_prob,
                                 use_one_hot_embeddings)
        param_dict = load_checkpoint(student_ckpt)
        new_param_dict = {}
        for key, value in param_dict.items():
            new_key = 'bert.' + key
            new_param_dict[new_key] = value
        load_param_into_net(self.bert, new_param_dict)

        # ----------------------
        # [need mindspore2.0]
        # ----------------------
        # compute head_importance and neuron_importance
        # head_importance: 12*12 matrix for head importance in 12 layers.
        # neuron_importance: list for neuron importance in 12 layers.

        # if training_phase == 'dynabertw':
        #     base_model, head_importance, neuron_importance = compute_neuron_head_importance\
        #         (self.bert, student_config, new_param_dict, eval_dataset=eval_dataset)
        #     reorder_neuron_head(base_model, head_importance, neuron_importance)

        self.cast = P.Cast()
        self.teacher_layers_num = teacher_config.num_hidden_layers
        self.student_layers_num = student_config.num_hidden_layers
        self.layers_per_block = int(self.teacher_layers_num / self.student_layers_num)
        self.is_rep_fit = is_rep_fit
        self.task_type = task_type
        self.temperature = temperature
        self.loss_mse = nn.MSELoss()
        self.lgt_fct = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
        self.select = P.Select()
        self.zeroslike = P.ZerosLike()
        self.dtype = student_config.dtype
        self.num_labels = num_labels
        self.soft_cross_entropy = SoftmaxCrossEntropy()
        self.compute_type = student_config.compute_type
        self.reshape = P.Reshape()
        self.training_phase = training_phase
        self.width_mult_list = [float(width) for width in "0.25, 0.5, 0.75, 1.0".split(',')]
        self.depth_mult_list = [1.]
        self.width_lambda1 = 1.0
        self.width_lambda2 = 0.1
        self.depth_lambda1 = 1.0
        self.depth_lambda2 = 1.0


    def construct(self,
                  input_ids,
                  input_mask,
                  token_type_id,
                  label_ids):
        """network with loss"""
        """
        1.Train DynaBERTw:
        teacher model: Fine-tuned BERT
        student model: Fine-tuned BERT ----> DynaBERTw

        2.Train DynaBERT:
        teacher model: DynaBERTw
        student model: DynaBERTw ----> DynaBERT
        """

        lossw = 0
        hidden_max_all, logits_max_all = [], []
        kept_layers_index = []
        if self.training_phase == 'dynabertw':
            teacher_logits_w, teacher_reps_w, _, _ = self.teacher(input_ids, token_type_id, input_mask)
        elif self.training_phase == 'dynabert':
            for width_mult in sorted(self.width_mult_list, reverse=True):
                teacher_logits_t, teacher_reps_t, _, _ = self.teacher(input_ids, token_type_id, input_mask)
                hidden_max_all.append(teacher_reps_t)
                logits_max_all.append(teacher_logits_t)

        # accumulate grads for all sub_networks
        for depth_mult in sorted(self.depth_mult_list, reverse=True):
            # model.apply set value
            self.bert.bert.apply(lambda m: setattr(m, 'depth_mult', depth_mult))

            # select teacher model layers for matching
            if self.training_phase == 'dynabert' or 'final_finetuning':
                n_layers = self.student_layers_num
                depth = round(depth_mult * n_layers)
                for i in range(depth):
                    kept_layers_index.append(math.floor(i / depth_mult))
                kept_layers_index.append(n_layers)

            # adjust width
            width_idx = 0
            for width_mult in sorted(self.width_mult_list, reverse=True):
                self.bert.bert.apply(lambda m: setattr(m, 'width_mult', width_mult))

                # stage 1: width-adaptive
                if self.training_phase == 'dynabertw':
                    # student model: 'dynabertw'
                    student_logits_w, student_reps_w, _, _ = self.bert(input_ids, token_type_id, input_mask)
                    if self.is_rep_fit:
                        selected_teacher_seq_output = ()
                        selected_student_seq_output = ()
                        for i in range(self.student_layers_num + 1):
                            selected_teacher_seq_output += (teacher_reps_w[i],)
                            selected_student_seq_output += (student_reps_w[i],)
                        rep_loss = 0
                        for i in range(self.student_layers_num + 1):
                            teacher_rep_w = selected_teacher_seq_output[i]
                            student_rep_w = selected_student_seq_output[i]
                            rep_loss += self.loss_mse(student_rep_w, teacher_rep_w)
                        lossw += rep_loss * self.width_lambda2
                    cls_loss = 0
                    if self.task_type == "classification":
                        cls_loss = self.soft_cross_entropy(student_logits_w / self.temperature,
                                                           teacher_logits_w / self.temperature)
                    elif self.task_type == "regression":
                        cls_loss = 0
                    lossw += cls_loss * self.width_lambda1


                # stage 2: width- and depth- adaptive
                elif self.training_phase == 'dynabert':
                    student_logits_t, student_reps_t, _, _ = self.bert(input_ids, token_type_id, input_mask)
                    cls_loss = 0
                    if self.task_type == "classification":
                        cls_loss = self.soft_cross_entropy(student_logits_t / self.temperature,
                                                           logits_max_all[width_idx] / self.temperature)
                    elif self.task_type == "regression":
                        cls_loss = 0
                    lossw += cls_loss * self.depth_lambda1
                    if self.is_rep_fit:
                        selected_teacher_seq_output = ()
                        selected_student_seq_output = ()
                        teacher_reps = list(hidden_max_all[width_idx][k] for k in kept_layers_index)
                        for i in range(self.student_layers_num + 1):  # 13层
                            selected_teacher_seq_output += (teacher_reps[i],)
                            selected_student_seq_output += (student_reps_t[i],)
                        rep_loss = 0
                        for i in range(self.student_layers_num + 1):
                            teacher_rep_t = selected_teacher_seq_output[i]
                            student_rep_t = selected_student_seq_output[i]
                            rep_loss += self.loss_mse(student_rep_t, teacher_rep_t)
                        lossw += rep_loss * self.depth_lambda2
                    width_idx += 1  # move to the next width

        print('lossw:', lossw)
        return self.cast(lossw, mstype.float32)


class BertTrainWithLossScaleCell(nn.Cell):
    """
    Encapsulation class of bert network training.

    Append an optimizer to the training network after that the construct
    function can be called to create the backward graph.

    Args:
        network (Cell): The training network. Note that loss function should have been added.
        optimizer (Optimizer): Optimizer for updating the weights.
        scale_update_cell (Cell): Cell to do the loss scale. Default: None.
    """
    def __init__(self, network, optimizer, scale_update_cell=None):
        super(BertTrainWithLossScaleCell, self).__init__(auto_prefix=False)
        self.network = network
        self.network.set_grad()
        self.weights = optimizer.parameters
        self.optimizer = optimizer
        self.grad = C.GradOperation(get_by_list=True,
                                    sens_param=True)
        self.reducer_flag = False
        self.allreduce = P.AllReduce()
        self.parallel_mode = context.get_auto_parallel_context("parallel_mode")
        if self.parallel_mode in [ParallelMode.DATA_PARALLEL, ParallelMode.HYBRID_PARALLEL]:
            self.reducer_flag = True
        self.grad_reducer = F.identity
        self.degree = 1
        if self.reducer_flag:
            self.degree = get_group_size()
            self.grad_reducer = DistributedGradReducer(optimizer.parameters, False, self.degree)
        self.is_distributed = (self.parallel_mode != ParallelMode.STAND_ALONE)
        self.cast = P.Cast()
        self.gpu_target = False
        if context.get_context("device_target") == "GPU":
            self.gpu_target = True
            self.float_status = P.FloatStatus()
            self.addn = P.AddN()
            self.reshape = P.Reshape()
        elif context.get_context("device_target") == "Ascend":
            self.alloc_status = P.NPUAllocFloatStatus()
            self.get_status = P.NPUGetFloatStatus()
            self.clear_before_grad = P.NPUClearFloatStatus()
        self.reduce_sum = P.ReduceSum(keep_dims=False)
        # self.depend_parameter_use = P.ControlDepend(depend_mode=1)
        self.base = Tensor(1, mstype.float32)
        self.less_equal = P.LessEqual()
        self.hyper_map = C.HyperMap()
        self.loss_scale = None
        self.loss_scaling_manager = scale_update_cell
        if scale_update_cell:
            self.loss_scale = Parameter(Tensor(scale_update_cell.get_loss_scale(), dtype=mstype.float32),
                                        name="loss_scale")

    @C.add_flags(has_effect=True)
    def construct(self,
                  input_ids,
                  input_mask,
                  token_type_id,
                  label_ids,
                  sens=None):
        """Defines the computation performed."""
        weights = self.weights
        loss = self.network(input_ids, input_mask, token_type_id, label_ids)
        if sens is None:
            scaling_sens = self.loss_scale
        else:
            scaling_sens = sens
        # alloc status and clear should be right before gradoperation
        if not self.gpu_target:
            init = self.alloc_status()
            self.clear_before_grad(init)
        grads = self.grad(self.network, weights)(input_ids,
                                                 input_mask,
                                                 token_type_id,
                                                 self.cast(scaling_sens,
                                                           mstype.float32))
        # apply grad reducer on grads
        grads = self.grad_reducer(grads)
        grads = self.hyper_map(F.partial(grad_scale, scaling_sens * self.degree), grads)
        grads = self.hyper_map(F.partial(clip_grad, GRADIENT_CLIP_TYPE, GRADIENT_CLIP_VALUE), grads)
        if not self.gpu_target:
            self.get_status(init)
            flag_sum = self.reduce_sum(init, (0,))
        else:
            flag_sum = self.hyper_map(F.partial(_grad_overflow), grads)
            flag_sum = self.addn(flag_sum)
            flag_sum = self.reshape(flag_sum, (()))

        if self.is_distributed:
            # sum overflow flag over devices
            flag_reduce = self.allreduce(flag_sum)
            cond = self.less_equal(self.base, flag_reduce)
        else:
            cond = self.less_equal(self.base, flag_sum)
        overflow = cond
        if sens is None:
            overflow = self.loss_scaling_manager(self.loss_scale, cond)
        if overflow:
            succ = False
        else:
            succ = self.optimizer(grads)
        ret = (loss, cond, scaling_sens)
        return F.depend(ret, succ)


    def construct(self,
                  input_ids,
                  input_mask,
                  token_type_id,
                  label_ids,
                  sens=None):
        """
        Args:
            input_ids (Tensor): the indices of input sequence tokens in the vocabulary.
            input_mask (Tensor): input sequence padding mask, where 0 indicates padding position.
            label_ids (Tensor): the indices of input sequence tokens in the vocabulary
        """

        weights = self.weights
        init = False
        loss = self.network(input_ids,
                            input_mask,
                            token_type_id,
                            label_ids)
        if sens is None:
            scaling_sens = self.loss_scale
        else:
            scaling_sens = sens

        if not self.gpu_target:
            init = self.alloc_status()
            init = F.depend(init, loss)
            clear_before_grad = self.clear_before_grad(init)
            scaling_sens = F.depend(scaling_sens, clear_before_grad)
        grads = self.grad(self.network, weights)(input_ids,
                                                 input_mask,
                                                 token_type_id,
                                                 label_ids,
                                                 self.cast(scaling_sens,
                                                           mstype.float32))
        grads = self.hyper_map(F.partial(grad_scale, scaling_sens), grads)
        grads = self.hyper_map(F.partial(clip_grad, GRADIENT_CLIP_TYPE, GRADIENT_CLIP_VALUE), grads)
        if self.reducer_flag:
            grads = self.grad_reducer(grads)
        if not self.gpu_target:
            init = F.depend(init, grads)
            flag = self.get_status(init)
            init = F.depend(init, flag)
            flag_sum = self.reduce_sum(init, (0,))
        else:
            flag_sum = self.hyper_map(F.partial(_grad_overflow), grads)
            flag_sum = self.addn(flag_sum)
            flag_sum = self.reshape(flag_sum, (()))
        if self.is_distributed:
            flag_reduce = self.allreduce(flag_sum)
            cond = self.less_equal(self.base, flag_reduce)
        else:
            cond = self.less_equal(self.base, flag_sum)
        overflow = cond
        if sens is None:
            overflow = self.loss_scaling_manager(self.loss_scale, cond)
        if overflow:
            succ = False
        else:
            succ = self.optimizer(grads)
        ret = (loss, cond, scaling_sens)
        return F.depend(ret, succ)


class BertTrainCell(nn.Cell):
    """
    Encapsulation class of bert network training.

    Append an opimizer to the training network after that the construct
    function can be called to create the backward graph.

    Args:
        network (Cell): The training network. Note that loss function should have been added.
        optimizer (Optimizer): Optimizer for updating the weights.
        sens (Number): The adjust parameter. Default: 1.0.
    """
    def __init__(self, network, optimizer, sens=1.0):
        super(BertTrainCell, self).__init__(auto_prefix=False)
        self.network = network
        self.network.set_grad()
        self.weights = optimizer.parameters
        self.optimizer = optimizer
        self.sens = sens
        self.grad = C.GradOperation(get_by_list=True,
                                    sens_param=True)
        self.reducer_flag = False
        self.parallel_mode = context.get_auto_parallel_context("parallel_mode")
        if self.parallel_mode in [ParallelMode.DATA_PARALLEL, ParallelMode.HYBRID_PARALLEL]:
            self.reducer_flag = True
        self.grad_reducer = F.identity
        self.degree = 1
        if self.reducer_flag:
            mean = context.get_auto_parallel_context("gradients_mean")
            self.degree = get_group_size()
            self.grad_reducer = DistributedGradReducer(optimizer.parameters, mean, self.degree)
        self.cast = P.Cast()
        self.hyper_map = C.HyperMap()

    def construct(self,
                  input_ids,
                  input_mask,
                  token_type_id,
                  label_ids):
        """Defines the computation performed."""
        weights = self.weights
        loss = self.network(input_ids,
                            input_mask,
                            token_type_id,
                            label_ids)
        grads = self.grad(self.network, weights)(input_ids,
                                                 input_mask,
                                                 token_type_id,
                                                 label_ids,
                                                 self.cast(F.tuple_to_array((self.sens,)),
                                                           mstype.float32))
        # apply grad reducer on grads
        grads = self.grad_reducer(grads)
        grads = self.hyper_map(F.partial(clip_grad, GRADIENT_CLIP_TYPE, GRADIENT_CLIP_VALUE), grads)
        succ = self.optimizer(grads)
        return F.depend(loss, succ)